#include "Registration.hpp"
#include "SubdivisionAlgorithm.hpp"
#include "PolyMesh.hpp"
#include "linearsystem.h"
#include <stdio.h>
#include "DebugHelper.hpp"

#include "boost/multi_array.hpp"
#include "boost/numeric/ublas/matrix.hpp"
#include "boost/numeric/ublas/matrix_proxy.hpp"
#include "boost/numeric/ublas/matrix_sparse.hpp"
#include "boost/numeric/ublas/operation.hpp"
#include "boost/numeric/ublas/operation_sparse.hpp"
#include "MeshAlgorithm.hpp"

namespace ublas = boost::numeric::ublas;
//typedef boost::numeric::ublas::mapped_matrix<float> SparseMatrix;
typedef boost::numeric::ublas::compressed_matrix<float> SparseMatrix;
typedef ublas::matrix<float> Matrix;

namespace
{
   const int cSUB_LEVEL = 3;
   const int tettris[4][3]={{0,1,2},{0,2,3},{0,3,1},{1,3,2}};
   const int tettip[4] =    {3,      1,      2,      0};
   const int octtris[8][3]={{0,1,2},{3,4,5},{0,4,3},{0,2,4},{2,5,4},{2,1,5},{1,3,5},{1,0,3}};
}

static float getVolume(Vec3& a, Vec3& b, Vec3& c, Vec3& d)
{
   Vec3 na = b - a;
   Vec3 nb = c - b;
   Vec3 nc = c - d;
   Vec3 n = cross(na,nb);
   return nc.dot(n)/6.0f;
}
void Registration::AlignLS2(PolyMesh&     aMesh,
                            MatPointList& aLandmarks,
                            float         fitWt,
                            float         energyWt,
                            int           subLevel,
                            int           iterations)
{
   //int i,j,k ;
   int maxIterations = 100 ;
   int fitSubLevel = 3 ;
   int energySubLevel = 2 ;

   FaceList& faces = *aMesh.crease_faces;
   VertexList& verts = *aMesh.verts;
   
   // Find interior vertices
   int* vertflag = new int[ verts.size() ] ;
   for (int i = 0 ; i < verts.size() ; i ++ )
   {
      vertflag[ i ] = 0 ;
   }
   for (int i = 0 ; i < faces.size() ; i ++ )
   {
      for (int j = 0 ; j < 3 ; j ++ )
      {
         vertflag[ faces[i].vertex[j] ] = 1 ;
      }
   }

   // allocate ata, atb, x and temporary vars
   printf("Allocating...\n") ;
   int numvar = verts.size();
   typedef boost::multi_array<float, 2> array_float2;

   // Subdivide mask
   printf("Subdividing...\n") ;
   float* maskl = new float[verts.size() * verts.size()];
   for (size_t i = 0; i < verts.size(); ++i)
   {
      int ind = verts.size() * i;
      for (size_t j = 0; j < verts.size(); ++j)
      {
         maskl[ind + j] = (j==i) ? 1 : 0;
      }
   }

   // Subdivide the mask
   PolyMesh subMaskMesh;
   subMaskMesh.clone(aMesh);
   maskl = SubdivisionAlgorithm::SubdivideMask(subMaskMesh, subLevel, maskl, verts.size());
   
   //int numIntFaces = (4 * aMesh.polys.size() - aMesh.crease_faces.size()) / 2;
   std::vector<MeshAlgorithm::TetPair> tetPairs;
   MeshAlgorithm::GetTetPairs(aMesh, tetPairs);
   int numIntFaces = tetPairs.size();
   int numrows = aLandmarks.size() + numIntFaces;

   printf("Adding fit terms...\n") ;
   //float** a = new float * [ numrows ] ;

   for (int iterNum = 0; iterNum < iterations; ++iterNum)
   {
      printf("Iteration %i\n", iterNum);
      PolyMesh subMesh;
      subMesh.clone(aMesh);
      SubdivisionAlgorithm::SubdivideMesh(subMesh, subLevel);

      SparseMatrix a(numrows, numvar);
      Matrix b(3, numrows);

      // collect list of vertices used in crease faces
      std::vector<int> creaseFaceVertices;
      for (size_t i = 0; i < subMesh.crease_faces.size(); ++i)
      {
         Face& f = subMesh.crease_faces[i];
         creaseFaceVertices.insert(creaseFaceVertices.end(), f.vertex.begin(), f.vertex.end());
      }
      std::sort(creaseFaceVertices.begin(), creaseFaceVertices.end());
      creaseFaceVertices.erase(std::unique(creaseFaceVertices.begin(), creaseFaceVertices.end()), creaseFaceVertices.end());

      MatPointList subMeshMatPoints;
      MeshAlgorithm::ExtractMatPoints(subMesh.GetCreaseMesh(), subMeshMatPoints);
      for (int i = 0 ; i < aLandmarks.size(); ++i)
      {
         // for each landmark, find closest mesh vertex
         MatPoint& landmark = aLandmarks[i];
         Vec3 lm = aLandmarks[i].point;
         int ind = -1 ;
         float mindis2 = 1000000.0f * 1000000.0f;
         //for (int j = 0 ; j < creaseFaceVertices.size(); ++j)
         for (size_t j = 0; j < subMeshMatPoints.size(); ++j)
         {
            // TODO: Only consider surface vertices [ dont need check? ]
            MatPoint& pt = subMeshMatPoints[j];
            if (pt.materials == landmark.materials)
            {
               int fi = j;
               float dis = (subMesh.verts[fi] - lm).length_squared();
               if ( dis < mindis2 )
               {
                  mindis2 = dis ;
                  ind = fi ;
               }
            }
         }

         if (ind >= 0)
         {
            // Get A
            for (int j = 0 ; j < numvar ; j ++ )
            {
               float val = maskl[verts.size() * ind + j] * fitWt;
               if (val != 0.0f)
               {
                  a(i,j) = val;
               }
            }

            // Get b for each x,y,z component
            for (int j = 0 ; j < 3 ; ++j)
            {
               float bVal = lm[ j ] * fitWt;
               if (bVal != 0.0f)
               {
                  b(j,i) = bVal;
               }
            }
         }
      }


      // Add energy terms
      printf("Adding energy terms...\n") ;

      // Find all interior faces
      int tettip[4]={3,1,2,0};

      int ind = aLandmarks.size() ;
      for (int i = 0 ; i < tetPairs.size() ; i ++ )
      {
         MeshAlgorithm::TetPair& p = tetPairs[i];

         // Get relevant vertices
         int ip1 = p.tetTips[0];
         int ip2 = p.tetTips[1];
         int im1 = p.verts[0];
         int im2 = p.verts[1];
         int im3  = p.verts[2];
         Vec3 p1 = verts[ip1];
         Vec3 p2 = verts[ip2];
         Vec3 m1 = verts[im1];
         Vec3 m2 = verts[im2];
         Vec3 m3 = verts[im3];

         // Get different volumes
         float vol1 = getVolume( m3,m2,m1,p2 ) ;
         float vol2 = getVolume( m1,m2,m3,p1 ) ;
         float vol3 = getVolume( p1,m3,m2,p2 ) ;
         float vol4 = getVolume( p1,m1,m3,p2 ) ;
         float vol5 = getVolume( p1,m2,m1,p2 ) ;
         float normvol = vol1 + vol2 ;
         normvol *= normvol ;
         if (normvol > 1.0E-12)
         {
            a(ind,ip1) = vol1 / normvol * energyWt ;
            a(ind,ip2) = vol2 / normvol * energyWt ;
            a(ind,im1) = -vol3 / normvol * energyWt ;
            a(ind,im2) = -vol4 / normvol * energyWt ;
            a(ind,im3) = -vol5 / normvol * energyWt ;
         }

         //if (!(math::is_finite(a[ind][ip1]) && 
         //math::is_finite(a[ind][ip2])&& 
         //math::is_finite(a[ind][im1])&& 
         //math::is_finite(a[ind][im2])&& 
         //math::is_finite(a[ind][im3])))
         //{
         //   int stophere=1;
         //}
         for (int j = 0 ; j < 3 ; j ++ )
         {
            b(j,ind) = 0;
         }
         ind ++ ;
      }


      printf("Adding to matrix...\n") ;
      DebugHelper::CheckMemory();
      SparseMatrix aT = ublas::trans(a);
      Matrix ata(aT.size1(), a.size1());
      ublas::axpy_prod(aT,a,ata);
      Matrix atb(b.size1(), a.size2());
      ublas::axpy_prod(b,a,atb);

      float** ataRows = new float*[ata.size1()];
      for (size_t i = 0; i < ata.size1(); ++i)
      {
         ataRows[i] = &ata(i,0); 
      }
      //std::cout << "ATB: " << atb.size1() << ',' << atb.size2() << '\n';
      float** atbRows = new float*[atb.size1()];
      for (size_t i = 0; i < atb.size1(); ++i)
      {
         atbRows[i] = &atb(i,0); 
      }


      // Solving
      printf("Solving...\n") ;
      std::vector<float> x(numrows);
      for (int j = 0 ; j < 3 ; j ++ )
      {
         for (int i = 0 ; i < verts.size() ; i ++ )
         {
            x[i] = verts[i][j];
         }
         LinearSystemSolver::solveD( numvar, ataRows, atbRows[j], &x[0], SOLVE_CONJUGATEGRADIENT ) ;
         float avg = 0 ;
         for (int i = 0 ; i < verts.size() ; i ++ )
         {
            verts[i][j] = x[i];
            avg += x[i] ;
         }
      }

      //delete[] aRows;
      //delete[] bRows;
      delete[] ataRows;
      delete[] atbRows;
   }

}
